import os
import pandas as pd
from WindPy import w
from tqdm import tqdm
from pymongo import MongoClient
from tjdutils.utils import TjdDate
from tjdutils.utils import output_dir


def get_data_from_edb(id_: str, name: str):
    print(id_, name)
    if not w.isconnected():
        w.start()
    data = w.edb(id_, beginTime='2002-01-01', endTime='2022-02-10')
    temp_dic = {
        '__id': id_,
        'df': pd.DataFrame(
            data.Data,
            columns=[pd.to_datetime(str(s)) for s in data.Times],
            # index=data.Codes
            index=[name]
        ).T,
        'name': name
    }
    return temp_dic


def get_data_from_mongo(id_li: str, name_li: str):
    coll = loc_zq_collection('wind_edb', 'data')
    data = list(coll.find({'__id': {'$in': id_li}}))
    assert len(data) > 0
    id_map = dict(zip(id_li, name_li))
    return {v['__id']: get_temp_dic(v, id_map) for v in data}


def get_temp_dic(v, id_map):
    id_ = v['__id']
    name = id_map[id_]
    print(id_, name)
    return {
        '__id': id_,
        'df': pd.DataFrame(v['data'], index=[pd.to_datetime(str(s)) for s in v['index']], columns=[name]),
        'name': name
    }


def get_features_by_basic_table(kwargs: dict, y_name):
    # 返回时间序列dict {'__id':,'df:,'name':}
    if kwargs['x_source'] == 'wind':
        feature_local_map = {k: v for k, v in kwargs['feature_map_wind'].items() if k in [y_name]}
        for k, v in feature_local_map.items():
            df_industry = pd.read_excel(v['x_source_name'], index_col=0)  # 读基础数据（提前规划好）
            map_dic = dict(zip(df_industry['codes'], df_industry['name']))  # id 与 name 的对应
        return {id_: get_data_from_edb(id_, name) for id_, name in map_dic.items()}
    elif kwargs['x_source'] == 'mongo':
        feature_map = {k: v for k, v in kwargs['feature_map_mongo'].items() if k in [y_name]}
        print('looking keys in', feature_map.keys(), 'from mongo @192.168.1.9')
        assert len(feature_map) > 0
        for k, v in feature_map.items():
            meta_coll = loc_zq_collection(v, 'metadata')
            metadata = list(meta_coll.find())
            id_li = [s['__id'] for s in metadata]
            name_li = [s['name'] for s in metadata]
            map_dic = get_data_from_mongo(id_li, name_li)
        return map_dic


def drop_inf(data, value=10):
    dt = data > value
    dt_s = dt.sum()
    dt_left = dt_s[dt_s > 0]
    names = dt_left.index.to_list()
    re = data.drop(names, axis=1)
    for i, name in enumerate(names):
        print(i, ">>", name)
    return re


def get_daily_yoy(data):
    data.index = [str(t)[:10] for t in data.index]
    concater = []
    for t in tqdm(data.index):
        t_all = TjdDate(t).dt
        temp_last_year = data.loc[t_all['lymfd']:t_all['lymld'], :]
        temp_this_year = data.loc[t_all['mfd']:t, :]
        mean_last_year = temp_last_year.mean(axis=0)
        mean_this_year = temp_this_year.mean(axis=0)
        mean_seires = mean_this_year / mean_last_year - 1
        mean_seires.name = t
        concater.append(mean_seires)
    re = pd.concat(concater, axis=1)
    re = re.T
    re_inter = re.interpolate()
    are = re_inter.fillna(method="ffill")
    are.index = pd.to_datetime(are.index)
    are = drop_inf(are)
    return are


def set_index(data):
    old_t = pd.to_datetime(data.index.to_list())
    t_li = pd.date_range(start=data.index[0], end=data.index[-1], freq='d')
    li = []
    for t in t_li:
        if t in old_t:
            se = data.loc[t, :]
        else:
            se = pd.Series(index=data.columns, name=t)
        li.append(se)
    new_data = pd.concat(li, axis=1)
    return new_data.T


def get_data(kwargs, y_name):
    re_dic = get_features_by_basic_table(kwargs, y_name)
    c = []
    for _, s in re_dic.items():
        dd = s['df']
        dd.sort_index(inplace=True)
        c.append(dd)
    df = pd.concat(c, axis=1)
    df.sort_index(inplace=True, ascending=True)
    return df


def ini_x_data(kwargs: dict, y_name: str):
    save_path = f'../input/{y_name}.pkl'
    print(save_path)
    if not os.path.exists(save_path) or kwargs['reset_x']:
        my_data = get_data(kwargs, y_name)
        my_data.index = pd.to_datetime(my_data.index)
        my_data = my_data[my_data.index > '2009-12-31']
        my_data = pd.DataFrame(my_data, dtype="float")
        my_data = set_index(my_data)
        my_data_inter = my_data.interpolate()
        my_data_fill = my_data_inter.fillna(method="ffill")
        yoy_data = get_daily_yoy(my_data_fill)
        output_dir()
        yoy_data.to_pickle(save_path)
        print(save_path, 'done!!')
    else:
        yoy_data = pd.read_pickle(save_path)
    return yoy_data


def loc_collection(db_name, collection_name):
    conn = MongoClient("mongodb://192.168.3.179:27017")
    coll_ = conn[db_name][collection_name]
    return coll_


def loc_zq_collection(db_name, collection_name):
    conn = MongoClient("mongodb://192.168.1.9:27017")
    coll_ = conn[db_name][collection_name]
    return coll_


def get_sw_names(level='L1'):
    coll_ = loc_zq_collection('wind_wsd', 'metadata')
    data = list(coll_.find({'level': level, }))
    id_ = [s['__id'] for s in data]
    names = [s['name'] for s in data]
    # re = list(loc_zq_collection('wind_wsd', 'data').find({"__id": {'$in': id_}}))
    # df_li = [pd.DataFrame(data=s['data'], index=s['index'], columns=s['columns']) for s in re]
    # df = pd.concat(df_li, axis=1)
    # df.columns = names
    return names


def cum_to_monthly_yoy(df):
    # 累计季度值变成当季值
    qoq_idx = df.index.to_list()
    qoq_value = pd.DataFrame(index=qoq_idx, columns=df.columns)
    for i, (quarter_t, quarter_se) in enumerate(df.iterrows()):
        k_last_quarter = qoq_idx[qoq_idx.index(quarter_t) - 1]
        if pd.to_datetime(quarter_t).month in [3]:
            qoq_se = df.loc[quarter_t, :]
        else:
            qoq_se = df.loc[quarter_t, :] - df.loc[k_last_quarter, :]
        qoq_value.loc[quarter_t, :] = qoq_se
    return qoq_value


def get_sw_data(level='L1', mode='营收'):
    mode_map = {'营收': 'sec_gr_sum_glb', '净利润': 'sec_np_sum_glb'}
    names = get_sw_names(level)
    data = list(loc_zq_collection('wind_wses', 'metadata').find({'name': {'$in': names, }, 'fields': mode_map[mode]}))
    id_s = [s['__id'] for s in data]
    ret = list(loc_zq_collection('wind_wses', 'data').find({'__id': {'$in': id_s, }}))
    df_li = [pd.DataFrame(data=s['data'], index=s['index'], columns=s['columns']) for s in ret]
    df = pd.concat(df_li, axis=1)
    names = [f"{s}_{mode}" for s in names]
    df.columns = names
    d = df.loc['2008-03-31':'2021-09-30', :]
    qoq = cum_to_monthly_yoy(d)
    qoq_yoy = qoq.pct_change(4).iloc[4:, :]
    return qoq_yoy


def ini_y_data(kwargs: dict):
    if kwargs['y_source'] == '申万数据':
        y_save_name = kwargs['y_save_name']
        reset_y = kwargs['reset_y']
        y_save_path = f'../input/{y_save_name}.xlsx'
        if not os.path.exists(y_save_path) or reset_y:
            y_data = get_sw_data()
            y_data.to_excel(y_save_path)
            print('y_data ini sucess!!!')
        else:
            y_data = pd.read_excel(y_save_path, index_col=0)
        return y_data


if __name__ == "__main__":
    print(666)
